# -*- coding: utf-8 -*-
# @Author  : longbhu
# @Time    : 2025/3/21 13:57
# @Function: 第二步：替换代码为新国标;然后合并成一个大图，用来再次裁剪。
import glob
import os

import rasterio
import numpy as np

def replace_values(tif_path, output_path):
    # 打开 TIFF 文件
    with rasterio.open(tif_path) as src:
        # 读取 TIFF 文件中的数据
        data = src.read(1)  # 假设 TIFF 文件只有一个波段

        # 获取元数据以保持相同的属性
        meta = src.meta

    # 定义对应关系
    mapping = {
        101: 1101,
        102: 1102,
        103: 1201,
        104: 1202,
        105: 1301,
        106: 2101,
        107: 2102,
        108: 2201,
        109: 5201,
        110: 5202,
        111: 5201,  # 注意：键是唯一的，这里的 111 会覆盖之前的 109 键对应的值
        112: 6205,
        21: 3100,
        22: 3200,
        23: 3400,
        24: 6206,
        31: 4101,
        32: 4102,
        33: 4103,
        34: 4201,
        35: 4202,
        36: 4301,
        37: 4302,
        41: 5101,
        42: 5102,
        51: 6101,
        52: 6102,
        53: 6103,
        54: 6104,
        61: 1401,
        62: 2301,
        63: 3400,  # 注意：键是唯一的，这里的 63 会覆盖之前的 23 键对应的值
        64: 3900,
        65: 8201,
        66: 8202,
        67: 7101,
        68: 7301,
        69: 8101
    }

    # 替换数据
    for old_value, new_value in mapping.items():
        data = data.astype(np.uint16)  # 更改为 uint16 类型
        data[data == old_value] = new_value

    # 更新元数据以匹配单个波段
    meta.update(
        compress='lzw',
        # count=1,
        dtype = np.uint16,
        nodata=255,
    )

    # 写入新的 TIFF 文件
    with rasterio.open(output_path, 'w', **meta) as dst:
        dst.write(data, 1)

def process_all_tifs(input_dir, output_dir):
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 获取所有 TIFF 文件的路径
    tif_files = glob.glob(os.path.join(input_dir, '*.tif'))

    for tif_path in tif_files:
        # 获取文件名
        file_name = os.path.basename(tif_path)
        # 构建输出路径
        output_path = os.path.join(output_dir, file_name)
        # 调用 replace_values 函数
        replace_values(tif_path, output_path)
        print(f"Processed {tif_path} and saved to {output_path}")

# 使用示例
input_dir = r'H:\land30m2010'
output_dir = r'H:\land30m2010\制作成新国标标准'
process_all_tifs(input_dir, output_dir)
